#include <math.h>
#include <string.h>
#include <ctype.h>

#include "base64.h"

static char base64_table[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
    "abcdefghijklmnopqrstuvwxyz" "0123456789+/";

int base64_needed_encoded_length(int length_of_data)
{
	int nb_base64_chars;
	nb_base64_chars = (length_of_data + 2) / 3 * 4;

	return nb_base64_chars + /** base64 char incl padding */
	    (nb_base64_chars - 1) / 76 +
	    1;			 /** NUL termination of string */
}

int base64_needed_decoded_length(int length_of_encoded_data)
{
	return (int)ceil(length_of_encoded_data * 3 / 4);
}

/**
  Encode a data as base64.

  Note: We assert that dst is pre-allocated to correct size.
        See base64_needed_encoded_length().
*/

int base64_encode(const void *src, size_t src_len, char *dst)
{
	const unsigned char *s = (const unsigned char *)src;
	size_t i = 0;
	size_t len = 0;

	for (; i < src_len; len += 4) {
		unsigned c;

		if (len == 76) {
			len = 0;
			*dst++ = '\n';
		}

		c = s[i++];
		c <<= 8;

		if (i < src_len)
			c += s[i];
		c <<= 8;
		i++;

		if (i < src_len)
			c += s[i];
		i++;

		*dst++ = base64_table[(c >> 18) & 0x3f];
		*dst++ = base64_table[(c >> 12) & 0x3f];

		if (i > (src_len + 1))
			*dst++ = '=';
		else
			*dst++ = base64_table[(c >> 6) & 0x3f];

		if (i > src_len)
			*dst++ = '=';
		else
			*dst++ = base64_table[(c >> 0) & 0x3f];
	}
	*dst = '\0';

	return 0;
}

static unsigned int pos(unsigned char c)
{
	return (unsigned int)(strchr(base64_table, c) - base64_table);
}

#define SKIP_SPACE(src, i, size) \
	do { \
		while (i < size && isspace(* src)) { \
			++i; ++src; \
		} \
		if (i == size) \
			break; \
	} while (0); \
	if (i >= size) \
		break

int base64_decode(const char *src_base, size_t len,
		  void *dst, const char **end_ptr)
{
	char b[3];
	size_t i = 0;
	char *dst_base = (char *)dst;
	char const *src = src_base;
	char *d = dst_base;
	size_t j;

	while (i < len) {
		unsigned c = 0;
		size_t mark = 0;

		SKIP_SPACE(src, i, len);

		c += pos(*src++);
		c <<= 6;
		if (++i >= len)
			break;

		SKIP_SPACE(src, i, len);

		c += pos(*src++);
		c <<= 6;
		if (++i >= len)
			break;

		SKIP_SPACE(src, i, len);

		if (*src != '=')
			c += pos(*src++);
		else {
			src += 2;
			/** There should be two bytes padding */
			i = len;
			mark = 2;
			c <<= 6;
			goto end;
		}
		c <<= 6;
		if (++i >= len)
			break;

		SKIP_SPACE(src, i, len);

		if (*src != '=')
			c += pos(*src++);
		else {
			src += 1;
			/** There should be one byte padding */
			i = len;
			mark = 1;
			goto end;
		}
		++i;

	end:
		b[0] = (c >> 16) & 0xff;
		b[1] = (c >> 8) & 0xff;
		b[2] = (c >> 0) & 0xff;

		for (j = 0; j < 3 - mark; j++)
			*d++ = b[j];
	}

	if (end_ptr != NULL)
		*end_ptr = src;

	/**
	 * The variable 'i' is set to 'len' when padding has been read, so it
	 * does not actually reflect the number of bytes read from 'src'.
	 */
	return i != len ? -1 : (int)(d - dst_base);
}

#ifdef BASE64_TESTCASE

#include <stdio.h>
#include <stdlib.h>
#include <assert.h>

int main(void)
{
	int i;
	size_t j;
	size_t k, l;
	size_t dst_len;
	size_t needed_length;

	for (i = 0; i < 500; i++) {
		/** Create source data */
		const size_t src_len = rand() % 1000 + 1;

		char *src = (char *)malloc(src_len);
		char *s = src;
		char *str;
		char *dst;

		assert(src);
		for (j = 0; j < src_len; j++) {
			char c = rand();
			*s++ = c;
		}

		/** Encode */
		needed_length = base64_needed_encoded_length(src_len);
		str = (char *)malloc(needed_length);
		assert(str);
		for (k = 0; k < needed_length; k++)
			str[k] = 0xff;
		    /** Fill memory to check correct NUL termination */
		assert(base64_encode(src, src_len, str) == 0);
		assert(needed_length == strlen(str) + 1);

		/** Decode */
		dst = (char *)malloc(base64_needed_decoded_length(strlen(str)));
		assert(dst);
		dst_len = base64_decode(str, strlen(str), dst, NULL);
		assert(dst_len == src_len);

		if (memcmp(src, dst, src_len) != 0) {
			printf("       --------- src ---------   --------- dst ---------\n");
			for (k = 0; k < src_len; k += 8) {
				printf("%.4x   ", (unsigned int)k);
				for (l = 0; l < 8 && k + l < src_len; l++) {
					unsigned char c = src[k + l];
					printf("%.2x ", (unsigned)c);
				}

				printf("  ");

				for (l = 0; l < 8 && k + l < dst_len; l++) {
					unsigned char c = dst[k + l];
					printf("%.2x ", (unsigned)c);
				}
				printf("\n");
			}
			printf("src length: %.8x, dst length: %.8x\n",
			       (unsigned int)src_len, (unsigned int)dst_len);
			assert(0);
		}
	}
	printf("Test succeeded.\n");
	return 0;
}

#endif /* BASE64_TESTCASE */
